#!/usr/bin/python
# coding=gbk
import sys,os,time,cPickle
from struct import unpack


class lzw:
	'''
	A class of LZW
	Useage:e=lzw()
	e.train(filename,type)
	e.test(filename)
	filename: file / dictory
	type: 'safe' or 'unsafe'
	'''
	
	def __init__(self,dicname='safe.dic',binary=1):
		self.binary=binary
		self.binarydict={'0':0,'b':1,'1':2,'CLR':3,'END':4}
		self.chardict={}
		d={}
		for i in range(0,256):		#dict in char
			d[chr(i)]=i
			d['CLR'],d['END']=256,257
		self.chardict=d
	
	def readfile(self,name='',mode='rb'):
		#try & except IOError
		f=file(name,mode)
		return f.read()
		
	def readdict(self,dicname):
		try:
			d=cPickle.load(file(dicname,'r'))
			print 'Load dic file "%s" done.' % dicname
			self.dic=d
		except:
			print 'Load dic file failed.\nCreate default dic.'
			self.dic={}
	
	def savedict(self,dicname):
		cPickle.dump(self.dic,file(dicname,'w'))
		print 'Save dic file "%s" done.' % dicname
	
	def cleardict(self):
		safedic=cPickle.load(file('safe.dic','r'))
		unsafedic=cPickle.load(file('unsafe.dic','r'))
		sharelist=[]
		for key in safedic:
			if key in unsafedic: sharelist.append(key)
		for key in sharelist:
			del safedic[key]
			del unsafedic[key]
		print 'Clear %s keys in dict' % len(sharelist)
		cPickle.dump(safedic,file('safe.dic','w'))
		cPickle.dump(unsafedic,file('unsafe.dic','w'))
	
	def train(self,file,type='safe'):
		files=[]
		if type(file)==type([]):
			files=file
		elif os.path.isdir(file):
			for i in os.listdir(file):
				files.append(file+'\\'+i)
		elif os.path.isfile(file):
				files.append(file)
		self.process(files,'train',type)
	
	def test(self,file):
		files=[];safenum=0;unsafenum=0
		if type(file)==type([]):
			files=file
		elif os.path.isdir(file):
			for i in os.listdir(file):
				files.append(file+'\\'+i)
		elif os.path.isfile(file):
				files.append(file)
		log_time=time.strftime('%Y-%m-%d %H_%M_%S')
		logtxt=time.strftime('%Y-%m-%d %H:%M:%S')+'\n'
		saferates=self.process(files,'test','safe')
		unsaferates=self.process(files,'test','unsafe')
		totalnum=len(saferates)
		for f in saferates:
			print '\nCompress rate with safe.dic:%2.1f%%' % (saferates[f]*100)
			print '\nCompress rate with unsafe.dic:%2.1f%%' % (unsaferates[f]*100)
			if saferates[f]>unsaferates[f]:
				print 'I think the file %s is safe.' % f
				safenum+=1
			else:
				print 'I think the file %s is not safe.' % f
				unsafenum+=1
		logtxt+='\ntesting %s files, safe=%s, unsafe=%s' % (totalnum,safenum,unsafenum)
		if '\\virus\\' in f:trueRate=unsafenum*1.0/totalnumm
		elif '\\nor\\' in f:truwRate=safenum*1.0/totalnum
		logtxt+='\ntrue rate: %2.1f' % trueRate
		f=open('test_'+log_time+'.txt','a')
		f.write(logtxt)
		f.close()

	def process(self,files,mode,type):
		times={};totalTime=0
		rates={}
		totalProcess=0;totalNumber=len(files)
		log_time=time.strftime('%Y-%m-%d %H_%M_%S')
		logtxt='Processing in %s mode,type=%s.\n' % (mode,type)
		logtxt+='Starttime:%s,totalNumber:%s\n\n' % (log_time,totalNumber)
		logtxt+='LocalTime, totalProcess, filename, duringTime, rate\n'
		# init dict
		dicname=type+'.dic'
		self.readdict(dicname)
		d=self.dic
		n=0
		#if(d=={}):
		#	if self.binary:
		#		#dict in binary code
		#		d=self.binarydict
		#		n=5
		#	else:
		#		d=self.chardict
		#		n=258
		#print 'Dict:',d	#[(i,d[i]) for i in sorted(d)]
		print 'Init Dic Done.len(dic)=%d' % len(d)

		for filename in files:
			try:
				startTime=time.time()
				print '-'*20
				print 'Processing %s in %s mode,type is %s.' % (filename,mode,type)
				
				s=''; ins=''; outs=''
				if(filename!=''):
					print 'Reading file %s ...' % filename
					s=self.readfile(filename)
					print 'len(%s)=%d byte' % (filename,len(s))
				
				print 'Unpacking to ascii...'
				l=unpack('%sB' % len(s), s)
				if(self.binary):
					print 'Coding to binary..'
					for i in l:
						ins+=bin(i)[2:]
				else:
					for i in l:
						ins+=chr(i)
				print 'ins done.len(ins)=%d' % len(ins)
				#print "input string:",ins
				
				# lzw coding
				#pre=ins[0]
				#for i in range(1,len(ins)):
				
				pre=''
				for i in range(len(ins)):
					suf=ins[i]
					if pre+suf in d:
						pre=str(d[pre+suf])
					else:
						outs+=pre;
						if(mode=='train'): d[pre+suf]=n;n+=1
						pre=suf
					#if(i%(len(ins)/4)==0): print "Done more 25%."

				rate=(len(ins)-len(outs))*1.0/len(ins)
				print "len(in)=%d; len(out)=%d; rate=%f" % (len(ins),len(outs),rate)
				print 'len(Dic)=%d' % len(d)
				#print 'Output string:',outs
				duringTime=time.time()-startTime
				print 'Processed using %.2f seconds.' % duringTime
				times[filename]=duringTime
				totalTime+=duringTime
				print 'Processed %s/%s,totaly using %.2f seconds.' % (len(times),len(files),totalTime)
				rates[filename]=rate
				totalProcess+=1
				logtxt+='%s, %s, %s, %s, %f\n'\
					% (time.strftime('%Y-%m-%d %H:%M:%S'),totalProcess,filename,duringTime,rate)
				if (totalProcess>0 and totalProcess%10==0):
					f=open('log_'+log_time+'.txt','a')
					f.write(logtxt)
					f.close()
					logtxt=''
			except:
				logtxt+='Error:'+filename+'\n'
		if(mode=='train'):
			self.dic = d
			self.savedict(dicname)
		if(mode=='test'):
			return rates	#,times
		logtxt+='\n\nProcess finished.\nTotalTime:%s mins %s seconds'\
			% (int(totalTime)/60,int(totalTime)%60)
		f=open('log_'+log_time+'.txt','a')
		f.write(logtxt)
		f.close()
		logtxt=''

if __name__=='__main__':
	lm=lzw()
	if len(sys.argv)>1:
		lm.test(sys.argv[1])
		#fs=[sys.argv[1]]
		#lm.process(fs,'train','dict')
	else:
		print lm.__doc__
	#raw_input('Press enter to exit.')